import numpy as np
import gradio as gr
from PIL import Image
import pinecone as pc
from tensorflow.keras.datasets import mnist

# Pinecone 配置
PINECONE_API_KEY = "pcsk_2WNR5R_95XGmZAVXXQYSc2txi8pJwXVWQaUg2WmdnT1nREKVJRSXdTH4CaJ3Lq9NcfKwD3"
INDEX_NAME = "developer-quickstart-py"
K_VALUE = 11

# 初始化 Pinecone 索引
pc_client = pc.Pinecone(api_key=PINECONE_API_KEY)
pinecone_index = pc_client.Index(INDEX_NAME)


def preprocess_sketchpad_image(image):
    """预处理手绘板图像（缩放为 8x8 并归一化）"""
    if image is None:
        return None
    # 从 Gradio Sketchpad 提取图像
    if isinstance(image, dict):
        image_array = image.get('image', None)
        if image_array is None:
            return None
        image = Image.fromarray(np.array(image_array, dtype=np.uint8))
    # 转为灰度图并缩放
    gray_img = image.convert('L')
    resized_img = gray_img.resize((8, 8), Image.Resampling.LANCZOS)
    inverted_img = Image.eval(resized_img, lambda x: 255 - x)  # 颜色反转
    np_img = np.array(inverted_img, dtype=np.float32) / 16  # 归一化到 0-16
    return np_img.flatten().tolist()


def predict_digit(image):
    """通过 Pinecone KNN 预测数字"""
    vector = preprocess_sketchpad_image(image)
    if vector is None:
        return None
    # 查询 Top-k 结果
    query_result = pinecone_index.query(
        vector=vector,
        top_k=K_VALUE,
        include_metadata=True
    )
    # 投票预测
    pred_labels = [int(match["metadata"]["label"]) for match in query_result["matches"]]
    return max(set(pred_labels), key=pred_labels.count)


# Gradio 界面
with gr.Blocks(title="Pinecone KNN 手写数字识别") as demo:
    gr.Markdown("# Pinecone KNN 手写数字识别")
    gr.Markdown("### 请在手写板绘制 0-9 的数字，点击提交查看预测结果")
    with gr.Row():
        sketchpad = gr.Sketchpad(height=200, width=200, label="手写板")
        result = gr.Number(label="预测结果", precision=0, interactive=False)
    with gr.Row():
        submit_btn = gr.Button("提交预测", variant="primary")
        clear_btn = gr.Button("清除手写板")
    # 绑定事件
    submit_btn.click(fn=predict_digit, inputs=sketchpad, outputs=result)
    clear_btn.click(fn=lambda: (None, None), inputs=[], outputs=[sketchpad, result])

if __name__ == "__main__":
    demo.launch(server_name="127.0.0.1", server_port=7862)